from __future__ import division, with_statement
"""
#===============================================================================
# License
#===============================================================================
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

#===============================================================================
Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Repository Path:    
$HeadURL: http://bryan-code.googlecode.com/svn/tags/PyJointTrack-0.1/JT/JT2.py $

Created on:         7/26/2009
Last Revised on:    $Date: 2009-08-04 02:38:40 +0000 (Tue, 04 Aug 2009) $
Version:            $Revision: 146 $
ID:                 $Id: JT2.py 146 2009-08-04 02:38:40Z bconrad2 $
Author:             $Author: bconrad2 $
#===============================================================================


#===============================================================================
# Documentation
#===============================================================================
A collection of components to perform image registration.

XRayCamera - a VTK camera that can be calibrated using JT parameters.
DRR - Generate 2D projection from 3D volume
ImageCanvas - a 2D image viewer
Registration - ITK components to perform 2D-2D registration (on projection and fixed image)
Volume - a VTK volume that has been setup with a texture mapper and custom transfer functions
Optimizer - OpenOpt
Metric - metrics that can be plugged into the ITK registraiton process
Frame - collection of models/volumes, fixed images, pose, calibration
"""

import JT

import pylab
import numpy as np

class Testing():
    """
    Functions related to the 2D3D Gold Standard Registration
    """
    def __init__(self):
        startPositionsPath = 'C:/Users/bryan/bryan-code/2D3D/startpositions_200_ordered.txt'
#===============================================================================
#    Read Start Positions
#===============================================================================
        f = open(startPositionsPath)
        # Start positions are 200, 6 element poses, plus error
        self.startPoses = np.zeros((201,7))
        
        self.roi = [[210,187],[277,370]]
    
        for ii, line in enumerate(f):
            self.startPoses[ii,:] = np.asarray(line.split(),'f')
            
    def GetStartPose(self,n=999):
        """
        Get Start Pose number n. n=999 will return all poses
        """
        if n==999:
            return self.startPoses[:,0:6]
        else:
            return self.startPoses[n,0:6]
    
    def GetStartPoseError(self,n=999):
        """
        Get Start Pose number n. n=999 will return all poses
        """
        if n==999:
            return self.startPoses[:,6]
        else:
            return self.startPoses[n,6]
    
    def GetDRR(self, windowName='DRR'):
        """
        Rerurn a fully stocked DRR generator ready to go out of the box!
        """
        print "Setting up a DRR instance"
        inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        calFile = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'      
        
#===============================================================================
#    Generate Lateral Projection
#===============================================================================
        print "Creating DRR..."
        drr = JT.DRR()
        drr.SetRenderWindowSize((512,512))
        
#===============================================================================
#    Calibration
#===============================================================================
        print "Performing Calibration..."
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFile)
        cal.LoadStaFile(staFile)
        
        drr.SetExternalCalibration(cal)
        
        volume=JT.Volume()
        volume.SetVolumeFilename(inputVolFileName)
        volume.SetOriginalTransform(cal._VolumeT)
        volume.SetBodyCenter((3.906, 2.17, -7.378))
        drr.AddVolume(volume)
        
        drr.Render()
        drr._renWin.SetWindowName(windowName)
        
        print "Done Preparing DRR!"
        return drr
    
    def GetReg(self):
        """
        Setup a registration for testing purposes
        """
        print "Setting up a registration instance"
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        inputVolFileName = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mha'
        staFile = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
        calFile = 'C:/Users/bryan/bryan-code/trunk/TestData/ext_cal1.txt'
        
#===============================================================================
#    Generate Lateral Projection
#===============================================================================
        print "Creating DRR..."
        drr = JT.DRR()
        drr.SetRenderWindowSize((512,512))
        
#===============================================================================
#    Calibration
#===============================================================================
        print "Performing Calibration..."
        cal = JT.ExternalCalibration()
        cal.LoadConsolidatedCalibration(calFile)
        cal.LoadStaFile(staFile)
        
        drr.SetExternalCalibration(cal)

        volume=JT.Volume()
        volume.SetVolumeFilename(inputVolFileName)
        volume.SetOriginalTransform(cal._VolumeT)
        volume.SetBodyCenter((3.906, 2.17, -7.378))
        drr.AddVolume(volume)

#===============================================================================
#    Fixed Image
#===============================================================================
        print "Loading Fixed Image..."
        fixedImage = JT.FixedImage()
        fixedImage.SetFileName(fixedImageName)
        
        print "Preparing Registration..."
        reg = JT.Registration()
        reg.SetFixedImage(fixedImage.GetImage(0))
        reg.SetMovingImageGenerator(drr)
        reg.SetRegionOfInterest([[210,187],[277,370]])
        reg.SetImageMetric('bpc')
        pose = (0,0,0,0,0,0)
        val = reg.GetValue(pose)
        print "Done Preparing Registration!"
        
        return reg
    
    def GetFixedImage(self,type='vtk'):
        """
        Return a fluro image as a vtk image or python array.  Default is vtk,
        if numpy array is desired, specify type as python
        """
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
        img = genericImageReader(fixedImageName)
        if type.lower() == 'python':
            img = vtk2numpy(img)
    
        return img
    
    def ConvertITKImageToArray(self,image):
        """
        For some reason, the GetFixedImage function cannot return a python array
        directly so this function is required???  BUG.
        """
        pyimg=itk.PyBuffer.IF2.GetArrayFromImage(image)
        return pyimg
    
    def GetTransform(self,pose):
        """
        Given a 6 element vector (Tx,Ty,Tx,Rx,Ry,Rz) return a transformation
        matrix.  (NOT shifted or concatenated with anything else)
        """
        Tx,Ty,Tz,Rx,Ry,Rz = pose
        
        # Tcs is TStartTransformation (body center to start pose)
        Tcs = np.eye(4)
        Tcs[0,0] = np.cos(Ry)*np.cos(Rz)
        Tcs[0,1] = -np.cos(Rx)*np.sin(Rz) + np.cos(Rz)*np.sin(Rx)*np.sin(Ry)
        Tcs[0,2] = np.sin(Rx)*np.sin(Rz) + np.cos(Rx)*np.cos(Rz)*np.sin(Ry)
        Tcs[0,3] = Tx
        Tcs[1,0] = np.cos(Ry)*np.sin(Rz)
        Tcs[1,1] = np.cos(Rx)*np.cos(Rz) + np.sin(Rx)*np.sin(Ry)*np.sin(Rz)
        Tcs[1,2] = -np.cos(Rz)*np.sin(Rx) + np.cos(Rx)*np.sin(Ry)*np.sin(Rz)
        Tcs[1,3] = Ty
        Tcs[2,0] = -np.sin(Ry)
        Tcs[2,1] = np.cos(Ry)*np.sin(Rx)
        Tcs[2,2] = np.cos(Rx)*np.cos(Ry)
        Tcs[2,3] = Tz
        
        return Tcs
    
    def GetTestImages(self):
        """
        Return ROI image region for fixed and moving images.
        """
        import cPickle
        fid = open('C:/Users/bryan/bryan-code/trunk/Images/fim.pickle')
        fim = cPickle.load(fid)
        fid.close()
        fid = open('C:/Users/bryan/bryan-code/trunk/Images/mim.pickle')
        mim = cPickle.load(fid)
        fid.close()
        return fim,mim
    
    
if __name__ == '__main__':
    testing = Testing()
    drr = testing.GetDRR('test')
    reg = testing.GetReg()
    